from glob import glob
import os.path as path
import os 
import json
from sklearn.model_selection import train_test_split


def get_data_dic():
    data_dic = {
    }
    return data_dic


def to_json(data, dest, json_name):
    os.makedirs(dest, exist_ok=True)
    with open(os.path.join(dest, json_name), 'w') as f:
        json.dump(data, f, indent=4)
        
 
def split_dataset(data_list=None):
    image_train, image_test = train_test_split(data_list, test_size=0.2, random_state=32) 
    image_train, image_val = train_test_split(image_train, test_size=0.2, random_state=32)

    return image_train, image_val, image_test


def get_split_data(image_list, cls):
    """
        cls: [lax, lvot, sax]
    """
    data_list = []
    
    key = 'full_sample_' + str(cls)
    
    for i in image_list:        
        anchor_path = i 
        data_dic = get_data_dic()
        
        # 读取 full_sample 所有数据路径
        # full_sample = [lax, lvot, sax]
        full_sample_path = glob(i + '/*')
        full_sample_path = sorted(full_sample_path)
        
        for ori_path in full_sample_path:
            if str(cls) in ori_path:
                full_sample_data = ori_path

        assert str(cls) in full_sample_data
        data_dic[key] = full_sample_data
        
        # 获取样本索引
        indice = path.split(i)[-1]
        data_dic['indice'] = indice
        
        data_list.append(data_dic)
    return data_list
      
        
def get_data_json(data_list, dest, json_name, cls):
    train_list, val_list, test_list = split_dataset(data_list)

    train_data = get_split_data(train_list, cls)
    val_data = get_split_data(val_list, cls)
    test_data = get_split_data(test_list, cls)

    data = {
        "train": list(train_data), 
        "val": list(val_data),
        "test": list(test_data)
    }
    
    to_json(data, dest, json_name)
    
    
if __name__ == "__main__":
    root = "/dev/raid/doctor/wyx/ChallengeData/home2/Raw_data/MICCAIChallenge2024/ChallengeData/MultiCoil/Cine/TrainingSet"

    full_sample = glob(root + '/FullSample/*')
    images_snap_shot = glob(root + '/ImgSnapshot/*')
    taks1 = glob(root + '/Mask_Task1/*')
    task2 = glob(root + '/Mask_Task2/*')

    dest = './'
    json_name = 'sax_ks.json'
    a = get_data_json(full_sample, dest, json_name, cls='sax')

    